#include "doca_dpa_unit.h"

extern "C"
{
    flexio_func_t feature_process;
    flexio_func_t sequence_preprocess;
}

extern struct flexio_app *feature_device; //这个和cmake中的program_lib_name对应

DOCA_LOG_REGISTER(MAIN);

DEFINE_int32(threads_host, 1, "host mem threads");
DEFINE_int32(arg_batch_size, 100, "batch_size");
DEFINE_int32(arg_feature_num, 1, "feature num");
DEFINE_int32(arg_ele_num, 5, "ele_per_feature");
DEFINE_string(device_type, "cpu", "running on which device");
DEFINE_string(op_type, "fillnull", "test single op_type");
std::vector<std::string> input_oplist = {
            "logit",
            "boxcox",
            "onehot",
            "bucketize",
            "sigrid_hash",
            "clamp_list",
            "ngram",
            "firstx",
            "mapid",
            "fillnull",
            "cast"
        };

std::map<std::string, int> local_map;

void local_map_init() {
    local_map.clear();
    for (int i = 0; i < (int)input_oplist.size(); ++i) {
        local_map[input_oplist[i]] = i;
    }
}

void check_addr(void * addr) {
    DOCA_LOG_INFO("host_addr's value 1 = %d, host_addr's value 2 = %d\n", ((int *)(addr))[0], ((int *)(addr))[1]);
}

void test_simple_functionality() {

    int thread_num = FLAGS_threads_host;

    DpaUnit dpaunit("mlx5_0", thread_num);
    DOCA_LOG_INFO("finish dpaunit init");

    dpaunit.context_init(feature_device);
    DOCA_LOG_INFO("finish context init");
    
    int host_mem_size = 1024 * 1024;
    dpaunit.log_init(0);
    DOCA_LOG_INFO("finish log init");

    dpaunit.command_queue_init();
    DOCA_LOG_INFO("finish command queue init");

    void * host_addr = dpaunit.prepare_host_memory(host_mem_size);
    void * test_addr = dpaunit.prepare_host_memory(host_mem_size);
    printf ("addr = %ld\n", (int64_t)test_addr);
    DOCA_LOG_INFO("finish host mem init");

    dpaunit.thread_init();
    DOCA_LOG_INFO("finish thread info init");

    check_addr((void*)((uint64_t)host_addr + 8));

    xmh::Reporter::StartReportThread();
    xmh::Timer time_calc("task execute time");
    
    dpaunit.submit_task(feature_process, host_addr, host_mem_size, LOGIT);
    DOCA_LOG_INFO("finish submit task");
    dpaunit.poll_thread();

    time_calc.end();

    DOCA_LOG_INFO("finish all task");
    check_addr((void*)((uint64_t)host_addr + 8));
    dpaunit.submit_task(feature_process, host_addr, host_mem_size, LOGIT);
    dpaunit.poll_thread();
    check_addr((void*)((uint64_t)host_addr + 8));

    xmh::Reporter::StopReportThread();
    xmh::Reporter::Report();
}

int check_times = 0;
std::vector<int64_t> correct_check_list[2];

template<typename T>
void check_status(T * unit) {
    auto meminfo = unit->mem_sequence[(int)(unit->mem_sequence).size() - 1];
    size_t total_size = meminfo.size;
    size_t check_length = (size_t)10000000 > total_size ? total_size : (size_t)10000000;
    for (int i = 0; i < (int)check_length; ++i) {
        correct_check_list[check_times].push_back( ( (int64_t *)(meminfo.mem_addr) )[i] );
    }

    printf ("%d times' check[", check_times);
    for (int i = 0; i < 10; ++i) {
        printf ("%ld, ", correct_check_list[check_times][i]);
    } printf ("]\n");

    check_times ++;
    
    if (check_times == 2) {
        for (int i = 0; i < (int)check_length; ++i) {
            if (correct_check_list[0][i] != correct_check_list[1][i]) {
                printf ("diff_happens!!!!\n");
                printf ("diff with pos %d, %ld != %ld\n", i, correct_check_list[0][i], correct_check_list[1][i]);
                return;
            }
        }
        printf ("all pos value correct\n");
    }
}

void test_sequence_operation() {
    int thread_num = FLAGS_threads_host;
    local_map_init();

    (void) thread_num;
    int sequence_len = FLAGS_arg_batch_size;
    int feature_size = FLAGS_arg_feature_num;
    int ele_num      = FLAGS_arg_ele_num;

    // std::vector<int> op_type_list = std::vector<int> {CLAMP_LIST, FILLNULL, SIGRID_HASH /*,CLAMP_LIST, NGRAM, FIRSTX*/};
    // std::vector<int> op_type_list = std::vector<int> {BOXCOX, LOGIT, BUCKETIZE, FILLNULL, SIGRID_HASH, CLAMP_LIST, NGRAM, FIRSTX};
    std::vector<int> op_type_list = std::vector<int> {FILLNULL};
    printf ("cur_type : %d\n", local_map[FLAGS_op_type]);
    op_type_list.push_back(local_map[FLAGS_op_type]);
    /* Dpa thread execute part start*/
    if (FLAGS_device_type != "dpu") {
        DpaUnit dpaunit("mlx5_0", thread_num);
        DOCA_LOG_INFO("finish dpaunit init");

        dpaunit.context_init(feature_device);
        DOCA_LOG_INFO("finish context init");
        
        dpaunit.log_init(0);
        DOCA_LOG_INFO("finish log init");

        dpaunit.command_queue_init();
        DOCA_LOG_INFO("finish command queue init");

        DpaPreProcessUnit DpaPrepare(&dpaunit);

        DpaPrepare.host_mem_init();
        DpaPrepare.create_sequence(op_type_list, feature_size, sequence_len, ele_num);
        DpaPrepare.map_init();
        DpaPrepare.thread_init();
        DpaPrepare.mem_init(0);
        DpaPrepare.print_mem();

        DOCA_LOG_INFO("finish preprocess_all_unit init");

        xmh::Timer time_calc("task execute time");
        
        DpaPrepare.task_execute(sequence_preprocess);

        time_calc.end();
        check_status(&DpaPrepare);
    }
    /* Dpa thread execute part end*/

    /* Cpu(dpu) thread execute part start*/
    if (FLAGS_device_type != "dpa") {
        CpuPreProcessUnit CpuPrepare(FLAGS_threads_host);
        DOCA_LOG_INFO("start cpu preprocess");
        CpuPrepare.host_mem_init();
        DOCA_LOG_INFO("finish host_mem init");
        CpuPrepare.create_sequence(op_type_list, feature_size, sequence_len, ele_num);
        DOCA_LOG_INFO("finish sequence init");
        CpuPrepare.map_init();
        DOCA_LOG_INFO("finish map init");
        CpuPrepare.thread_init();
        DOCA_LOG_INFO("finish thread init");
        CpuPrepare.mem_init(0);
        DOCA_LOG_INFO("finish mem init");
        CpuPrepare.print_mem();

        DOCA_LOG_INFO("finish cpu_all_unit init");

        xmh::Timer cpu_time_calc("cpu task execute time");
        
        CpuPrepare.task_execute(sequence_preprocess);

        cpu_time_calc.end();
        check_status(&CpuPrepare);
        std::cout << "OpenMP最大线程数: " << omp_get_max_threads() << std::endl;
    }
    /* Cpu(dpu) thread execute part end*/
}

int main(int argc, char **argv) {

    // doca_error_t result;
	// struct doca_log_backend *sdk_log;
    // result = doca_log_backend_create_standard();
	// result = doca_log_backend_create_with_file_sdk(stderr, &sdk_log);
	// result = doca_log_backend_set_sdk_level(sdk_log, DOCA_LOG_LEVEL_INFO);
    // if (result != DOCA_SUCCESS) printf ("error register log\n");
    gflags::ParseCommandLineFlags(&argc, &argv, true);

    // struct doca_log_backend *stdout_logger = nullptr;
    // Assert(doca_log_backend_create_with_file(stdout, &stdout_logger) == DOCA_SUCCESS);

	// DOCA_LOG_INFO("Starting the sample");
    // test_simple_functionality();
    xmh::Reporter::StartReportThread();

    test_sequence_operation();

    xmh::Reporter::StopReportThread();
    xmh::Reporter::Report();
    
    return EXIT_SUCCESS;
}